from typing import Type

import aiofiles
from langchain_community.tools import ReadFileTool
from langchain_community.tools.file_management.read import ReadFileInput
from pydantic import BaseModel, ConfigDict

from app.config import config
from app.documentloader.pdf import load as load_pdf, extract_content, extract_page_content
from app.documentloader.word import load as load_word
from app.logger import logger
from app.tool.file.file_operator import LocalFileOperator, SandboxFileOperator


class ReadFile(ReadFileTool):
    name: str = "read_file"
    args_schema: Type[BaseModel] = ReadFileInput
    description: str = "Read file using local or sandbox file operators."

    model_config = ConfigDict(extra="allow")

    def __init__(self):
        super().__init__()
        self.operator = SandboxFileOperator() if str(
            config.get("sandbox", "use_sandbox")).lower() == 'true' else LocalFileOperator()

    async def is_pdf(self, path: str) -> bool:
        """Check if file is a PDF by extension and header."""
        if not str(path).lower().endswith('.pdf'):
            return False

        try:
            async with aiofiles.open(path, 'rb') as f:
                header = await f.read(4)
                return header == b'%PDF'
        except Exception as e:
            logger.error(f"Failed to read {path}: {str(e)}")
            return False

    async def is_word(self, path: str) -> bool:
        """Check if file is a Word document by extension and header."""
        # 检查文件扩展名
        word_extensions = {'.doc', '.docx'}
        if not any(str(path).lower().endswith(ext) for ext in word_extensions):
            return False

        try:
            async with aiofiles.open(path, 'rb') as f:
                header = await f.read(8)  # 读取前8个字节用于检查

                # DOC文件标识
                if header.startswith(b'\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1'):
                    return True

                # DOCX文件实际上是ZIP格式，检查是否有Word相关标识
                if header.startswith(b'PK\x03\x04'):  # ZIP文件头
                    # 可以进一步检查是否包含Word特定结构
                    # 但通常扩展名+ZIP头就足够判断DOCX了
                    return str(path).lower().endswith('.docx')

                return False
        except Exception as e:
            logger.error(f"Failed to read {path}: {str(e)}")
            return False

    async def read_file(self, path: str) -> str:
        """Read content from a local file."""
        if await self.is_pdf(path):
            return await extract_content(await extract_page_content(await load_pdf(path)))

        if await self.is_word(path):
            return await extract_content(await extract_page_content(await load_word(path)))

        logger.debug(f"File operator: {self.operator.__class__.__name__}")
        return await self.operator.read_file(path)

    def _run(self, file_path: str) -> str:
        import asyncio
        try:
            return asyncio.run(self.read_file(file_path))
        except Exception as e:
            logger.error(f"Error reading file: {str(e)}")
            return f"Error reading file: {str(e)}"
